#!/usr/bin/python

#Programmer: Aaron Steele
#Date: 1/22/2013
#Purpose: Splice together min-max values with their appropriate place in the pfam
#		  alignment.  Computes for 1 single protein family.
#Usage: pfam-mm-align.py <min-max file> <alignment file>

import sys
import string
import os
import pfam_mm_align as mm
 
if __name__ == "__main__":
	
	#Get the start dir
	start_dir = sys.argv[1]
	
	#Dive into the directory
	for dirname, dirnames, filenames in os.walk(start_dir):

		#For each pfam-directory look at the alignment file and mm file
		for subdirname in dirnames:
			fam_name = subdirname.replace("/","")
			print("Aligning mm for " + fam_name)

			#Get the mm and align file
			mm_file = start_dir +  subdirname + "/" + fam_name + ".mm"
			align_file = start_dir + subdirname +"/" + fam_name + ".align"

			mm.align_pfam(mm_file,align_file) 

def align_pfam(mm_file, align_file):
	#Open the output file
	out_file = mm_file + ".align"

	#Read in all ids and sequences from alignment file
	align_ids = list()
	align_values = list()
	infile = open(align_file, "r")

	line = infile.readline()
	while line:
		line = line.rstrip()

		#If line starts with ">" then it is id, otherwise it is alignment
		if line.startswith(">"):
			line = line.replace(">", "")
			align_ids.append(line)
		else:
			align_values.append(line)
			
		line = infile.readline()

	infile.close()

	#Read in all ids and sequences from mm file
	mm_ids = list()
	mm_values = list()
	infile = open(mm_file, "r")

	line = infile.readline()
	while line:
		line = line.rstrip()

		#If it starts with ">" then it is a id, otherwise mm values
		if line.startswith(">"):
			line = line.replace(">","")
			#Remove the organism id after the ":"
			ssplit = line.split(":")
			mm_ids.append(ssplit[0])
		else:
			if line.endswith(","):
				mm_values.append(line[:-1])
			else:
				mm_values.append(line)
		line = infile.readline()	
	infile.close()

	#Open the output file for 
	outfile = open(out_file,"w")

	#Process the ids finding matches
	for i in range(len(align_ids)):
		value_string = ""

		#Get the alignment and trip the first and last non "-" 8 codons
		alignment_id = align_ids[i]
		alignment = align_values[i]
		temp = alignment.replace("-","")

		#Proceed if the sequences is long enough
		if len(temp) > 48:
			#Trim the first 8 codons used in mm calculations
			trim_count = 0
			while trim_count < 24:
				if not alignment[0] == "-":
					trim_count += 1
				alignment = alignment[1:]

			#Trim the last 8 codons used in mm calculations
			trim_count = 0
			while trim_count < 24:
				if not alignment[-1] == "-":
					trim_count += 1
				alignment = alignment[:-1]

			outfile.write(">" + alignment_id + "\n")

			#Find the matching mm id
			for j in range(len(mm_ids)):
				if mm_ids[j] == alignment_id:
					value_string = mm_values[j]
					break

			#Replace alignment with mm-values
			values = value_string.split(",")
			temp = alignment.replace("-","")
			#print(alignment_id + ": <" + str(len(temp)/3) + ", " + str(len(values)) + ">")


			mm_index = 0
			for j in range(len(alignment)/3):
			
				#If a "-" then write a zero
				if alignment[3*j] == "-":
					outfile.write("0")			
							
				#Write the next mm_value
				else:
					outfile.write(values[mm_index])
					mm_index = mm_index + 1
			
				#print alignment[3*j] + ", " + str(j) + ", " + str(mm_index)	
				#Write a comma if not the last mm
				if mm_index < len(values)-1:
					outfile.write(",")

			#All done writting mm, time to print a newline
			outfile.write("\n")


	#All done, close the outputfiles
	outfile.close()


